/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <folly/SocketAddress.h>
#include <folly/io/async/EventBaseThread.h>
#include "presto_cpp/main/http/HttpClient.h"
#include "velox/common/memory/Memory.h"
#include "velox/exec/tests/utils/ReferenceQueryRunner.h"
#include "velox/vector/ComplexVector.h"

namespace facebook::presto::test {

class PrestoQueryRunner : public velox::exec::test::ReferenceQueryRunner {
 public:
  /// @param coordinatorUri Presto REST API endpoint, e.g. http://127.0.0.1:8080
  /// @param user Username to use in X-Presto-User header.
  PrestoQueryRunner(const std::string& coordinatorUri, const std::string& user);

  /// Converts Velox query plan to Presto SQL. Supports Values -> Aggregation or
  /// Window with an optional Project on top.
  ///
  /// Values node is converted into reading from 'tmp' table.
  ///
  /// @return std::nullopt if Values node uses types not supported by DWRF file
  /// format (DATE, INTERVAL, UNKNOWN).
  std::optional<std::string> toSql(
      const velox::core::PlanNodePtr& plan) override;

  /// Creates 'tmp' table using specified data, executes SQL query generated by
  /// 'toSql' and returns the results.
  ///
  /// @param sql SQL generated by 'toSql' method.
  /// @param input Data used in the Values node in the plan passed to 'toSql'
  /// method.
  /// @param resultType Expected type of the results.
  /// @return Data received from Presto.
  std::multiset<std::vector<velox::variant>> execute(
      const std::string& sql,
      const std::vector<velox::RowVectorPtr>& input,
      const velox::RowTypePtr& resultType) override;

  /// Executes Presto SQL query and returns the results. Tables referenced by
  /// the query must already exist.
  std::vector<velox::RowVectorPtr> execute(const std::string& sql);

 private:
  velox::memory::MemoryPool* rootPool() {
    return rootPool_.get();
  }

  velox::memory::MemoryPool* pool() {
    return pool_.get();
  }

  std::optional<std::string> toSql(
      const std::shared_ptr<const velox::core::AggregationNode>&
          aggregationNode);

  std::optional<std::string> toSql(
      const std::shared_ptr<const velox::core::WindowNode>& windowNode);

  std::optional<std::string> toSql(
      const std::shared_ptr<const velox::core::ProjectNode>& projectNode);

  std::string startQuery(const std::string& sql, http::HttpClient& client);

  std::string fetchNext(const std::string& nextUri, http::HttpClient& client);

  const folly::SocketAddress coordinatorUri_;
  const std::string user_;
  folly::EventBaseThread eventBaseThread_{false};
  std::shared_ptr<velox::memory::MemoryPool> rootPool_{
      velox::memory::defaultMemoryManager().addRootPool()};
  std::shared_ptr<velox::memory::MemoryPool> pool_{
      rootPool_->addLeafChild("leaf")};
};

} // namespace facebook::presto::test
